import psycopg2
from datetime import datetime
import os

# 获取脚本所在路径
curPath = os.path.split(os.path.realpath(__file__))[0]


def find_index(lst, target):
    for i, item in enumerate(lst):
        if item == target:
            return i
    return -1


def read_and_delete_column(file_path):
    import csv

    with open('{}/data/{}.csv'.format(curPath, file_path), 'r', newline='', encoding='utf-8') as csvfile:
        reader = csv.reader(csvfile)
        datas = [row for row in reader]

    # 替换空字符串为 None
    for row in datas:
        row.pop(0)  # 删除id 避免id重复报错
        for idx, value in enumerate(row):
            if value == "":
                row[idx] = None

    return datas


def get_schema(cursor, table):
    if table.find('.') != -1:
        table = table.split('.')[1]
    # 获取表的 schema 信息
    print(table)
    cursor.execute(f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table}'")
    schema = {row[0]: row[1] for row in cursor.fetchall()}
    # 删除 id 字段并返回
    schema.pop('id')

    return schema


class PostgreSQLUtil:
    def __init__(self, host, port, dbname, username, password):
        self.host = host
        self.port = port
        self.dbname = dbname
        self.username = username
        self.password = password
        self.conn = None

    def connect_to_db(self):
        try:
            self.conn = psycopg2.connect(host=self.host, port=self.port, dbname=self.dbname, user=self.username,
                                         password=self.password)
            print("Connected to the database successfully!")
        except Exception as e:
            print(f"Error connecting to the database: {e}")

    def close_conn(self):
        if self.conn:
            self.conn.close()
            print("Connection closed.")
        else:
            print("No connection to close.")

    def write_data(self, table, datas, time_field):
        if not self.conn:
            print("Please connect to the database first.")
            return
        cursor = self.conn.cursor()
        schema = get_schema(cursor, table)
        fields = list(schema.keys())
        # 将数据中的时间改为当前时间
        i = find_index(fields, time_field)
        for row in datas:
            row[i] = datetime.now()

        insert_query = f"INSERT INTO {table} ({', '.join(fields)}) VALUES ({', '.join(['%s' for _ in range(len(fields))])})"
        try:
            cursor.executemany(insert_query, data)
            self.conn.commit()
            print(f"Data written to {table} successfully!")
        except Exception as e:
            print(f"Error writing data to {table}: {e}")
        finally:
            cursor.close()


if __name__ == "__main__":
    pg_table_dict = {
        # 'pg_alarm_msg': ['create_time', 'td_alarm_msg'],
        'pg_assets_manage': ['create_time', 'tr_assets_manage'],
        # 'pg_content_audit_result': ['acquisition_time', 'td_content_audit_result'],
        'pg_deal_commond_req': ['create_time', 'command.td_deal_commond_req'],
        'pg_device_upgrade': ['create_time', 'td_device_upgrade'],
        'pg_evil_resource_monitor_req': ['create_time', 'command.td_evil_resource_monitor_req'],
        'pg_evil_packet_monitor_req': ['create_time', 'command.td_evil_packet_monitor_req'],
        'pg_evil_file_monitor_req': ['create_time', 'command.td_evil_file_monitor_req'],
        'pg_his_minitor_req': ['create_time', 'command.td_his_minitor_req'],
        'pg_his_evil_sample_req': ['create_time', 'command.td_his_evil_sample_req'],
        'pg_stock_req': ['create_time', 'command.td_stock_req'],
        # 'pg_monitor_reported': ['create_time', 'td_monitor_reported'],
        # 'pg_library_upgrade': ['create_time', 'td_library_upgrade'],
        # 'pg_content_audit': ['create_time', 'td_content_audit'],
        # 一次导入
        'pg_device_manage': ['create_time', 'tr_device_manage'],
        'pg_device_group': ['create_time', 'tr_device_group']
    }

    util_111 = PostgreSQLUtil(host="10.10.3.111", port=5432, dbname="network_security",
                          username="postgres", password="postgres")
    util_111.connect_to_db()
    for key, val in pg_table_dict.items():
        data = read_and_delete_column(key)
        util_111.write_data(val[1], data, val[0])

    util_111.close_conn()
